# -*- coding: UTF-8 -*-
from typing import List
from leetcode_helper import TestcaseHelper, SolutionBase


class Stack:
    def __init__(self):
        self.data = []

    def push(self, item):
        self.data.append(item)

    def pop(self):
        if self.is_empty():
            return None
        else:
            return self.data.pop()

    def size(self):
        return len(self.data)

    def is_empty(self):
        return self.size() == 0

    def top(self):
        if self.is_empty():
            return None
        else:
            return self.data[-1]

    def get_data_list(self):
        return self.data


class Solution(SolutionBase):
    # 解题的方法，命名和Leetcode的保持一致即可
    def countOfAtoms(self, formula: str) -> str:
        def handle_pop_stack(atom_stack: Stack, atom_number_str: str) -> bool:
            if atom_stack.size() == 0:
                return False

            if atom_number_str == '':
                atom_number_str = '1'

            atom_number = int(atom_number_str)
            temp_atoms = []
            top = atom_stack.top()
            if top[0] == ')':
                while atom_stack.size() > 0:
                    pair = atom_stack.pop()
                    if pair[0] == ')':
                        continue
                    elif pair[0] == '(':
                        for item in temp_atoms:
                            atom_stack.push(item)
                        return True
                    else:
                        new_pair = (pair[0], pair[1] * atom_number)
                        temp_atoms.append(new_pair)
                return True
            else:
                return False

        # isupper, 是否为大写
        # isdecimal, 是否为数字
        atom = ''
        number_str = ''
        stack = Stack()

        for i in range(len(formula)):
            ch = formula[i]
            if ch.isupper():
                if handle_pop_stack(stack, number_str) is True:
                    number_str = ''
                    atom = ''

                if atom == '':
                    atom = ch
                else:
                    if number_str == '':
                        number_str = '1'
                    stack.push((atom, int(number_str)))
                    atom = ch
                    number_str = ''
            elif ch.isdecimal():
                number_str += ch
            elif ch == '(':
                if handle_pop_stack(stack, number_str) is True:
                    number_str = ''
                    atom = ''

                if atom != '':
                    if number_str == '':
                        number_str = '1'
                    stack.push((atom, int(number_str)))
                    atom = ''
                    number_str = ''
                stack.push(('(', 0))
            elif ch == ')':
                if handle_pop_stack(stack, number_str) is True:
                    number_str = ''
                    atom = ''
                if atom != '':
                    if number_str == '':
                        number_str = '1'
                    stack.push((atom, int(number_str)))
                    atom = ''
                    number_str = ''
                stack.push((')', 0))
            else:
                atom += ch

        if atom != '':
            if number_str == '':
                number_str = '1'
            stack.push((atom, int(number_str)))
        else:
            if number_str != '':
                handle_pop_stack(stack, int(number_str))

        atom_dict = {}
        while stack.size():
            atom = stack.pop()
            if atom[0] == '(' or atom[0] == ')':
                continue
            atom_dict[atom[0]] = atom_dict.get(atom[0], 0) + atom[1]

        atoms = atom_dict.keys()
        atoms = sorted(list(atoms))

        ans = ''
        for key in atoms:
            number = atom_dict[key]
            if number == 1:
                ans += key
            else:
                ans += key
                ans += str(number)

        return ans


if __name__ == '__main__':
    solution = Solution()

    # 按照leetcode的题目说明，将基本的test case填入
    # 每条case的最后一个元素是正确的结果，即期待的结果
    # 使用元组来保存输入的数据，期待的结果
    testcases = [
        ('(H)', 'H'),
        ("((HHe28Be26He)9)34", "Be7956H306He8874"),
        ("NH", "HN"),
        ("Mg(OH)2", "H2MgO2"),
        ("H2O", "H2O"),
        ("K4(ON(SO3)2)2", "K4N2O14S4")
    ]

    for case_item in testcases:
        # 获取测试用的case和期待的执行结果
        input1 = case_item[0]
        expect = case_item[len(case_item) - 1]
        exec_result = solution.countOfAtoms(input1)

        # 判断执行结果，输出提示信息
        check_result = solution.check_result(expect, exec_result)
        TestcaseHelper.print_case(check_result, case_item, exec_result)

